#include <bits/stdc++.h>
using namespace std;

/*
498. 对角线遍历
已解答
中等
相关标签
premium lock icon
相关企业
给你一个大小为 m x n 的矩阵 mat ，请以对角线遍历的顺序，用一个数组返回这个矩阵中的所有元素。

 

示例 1：


输入：mat = [[1,2,3],[4,5,6],[7,8,9]]
输出：[1,2,4,7,5,3,6,8,9]
示例 2：

输入：mat = [[1,2],[3,4]]
输出：[1,2,3,4]
 

提示：

m == mat.length
n == mat[i].length
1 <= m, n <= 104
1 <= m * n <= 104
-105 <= mat[i][j] <= 105
*/

// 法一
class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        // 考虑特殊情况
        if (mat.empty() || mat[0].empty())  return {};

        int m = mat.size();      // 矩阵行数
        int n = mat[0].size();   // 矩阵列数
        vector<int> ans;      // 存储结果

        // 遍历所有索引和 0 --> m + n - 2
        for (int i = 0; i <= m + n - 2; ++i) {
            // 对角线的行索引范围
            int start = max(0, i - (n - 1));
            int end = min(m - 1, i);
            // 奇偶分开
            if (i % 2 == 0) {
                for (int j = end; j >= start; --j) {
                    ans.push_back(mat[j][i- j]);
                }
            } else {
                for (int j = start; j <= end; ++j) {
                    ans.push_back(mat[j][i - j]);
                }
            }
        }
        return ans;
    }
};

// 法二
// 做点预处理提升性能
class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        if (mat.empty() || mat[0].empty()) return {};

        int m = mat.size();
        int n = mat[0].size();
        vector<int> ans;
        ans.reserve(m * n);          // 预分配空间，提升效率

        // 对角线的编号 d = i + j，范围是 [0, m+n-2]
        for (int d = 0; d <= m + n - 2; ++d) {
            // 当前对角线的起始行、列
            int rowStart = max(0, d - n + 1);          // i 最小值
            int rowEnd   = min(d, m - 1);              // i 最大值

            if (d % 2 == 0) { // 偶数 d —— 右上方向
                for (int i = rowEnd; i >= rowStart; --i) {
                    int j = d - i;
                    ans.push_back(mat[i][j]);
                }
            } else {          // 奇数 d —— 左下方向
                for (int i = rowStart; i <= rowEnd; ++i) {
                    int j = d - i;
                    ans.push_back(mat[i][j]);
                }
            }
        }
        return ans;
    }
};